import functools
import importlib
import pickle
from concurrent.futures import ProcessPoolExecutor as pool
from enum import Enum

import torch
from scipy.stats import ranksums, wilcoxon

from config import device, thread_pool_size


class MessageType(Enum):
    INIT_RENDER_MSG = '0'
    MULTIPLE_TASK_EXECUTE = '1'
    MULTIPLE_TASK_MSG = '2'
    MULTI_TASK_INDICATORS = '3'
    MULTI_TASK_COMPLETED = '4'
    SIMPLE_TASK_RECEIVE = '5'
    SIMPLE_TASK_MSG = '6'
    SIMPLE_TASK_START = '7'
    SIMPLE_TASK_COMPLETED = '8'
    SIMPLE_TASK_INDICATORS = '9'
    MULTI_STATISTICAL_RESULT = '10'
    MULTI_STATISTICAL_CALC = '11'
    MULTI_DATA_2_EXCEL = '12'
    MULTI_DATA_2_EXCEL_COMPLETE = '13'
    MULTI_Drawing = '14'
    MULTI_Drawing_COMPLETE = '15'


class MsgContext:
    executor = pool(max_workers=thread_pool_size)
    single_populations = []
    simple_problem_objs = []
    multiple_populations = None
    multiple_problem_objs = None
    multiple_indicators = {}
    multiple_FE = None
    callback = None
    indicators_tensor = {}
    file_path_list = None
    tensor_size = None
    indicator_value_list = None
    num_of_run = 'num of run'
    statistical_func = {"ranksums": ranksums, "signrank": wilcoxon}

    @classmethod
    def calc_indicator(cls, indicator_data, intermediate_pops, problem_index, algorithm_index, nth_run):
        indicator_class = cls.get_cls_by_str(indicator_data['cls'])
        indicator_value = indicator_class.get_value(intermediate_pops[-1], cls.multiple_problem_objs[problem_index])
        # 将当前计算的指标值保存到cls.multiple_indicators中
        temp_indicator = cls.multiple_indicators.get(indicator_data['name'],
                                                     torch.zeros(cls.tensor_size, device=device, dtype=torch.double))
        temp_indicator[problem_index][algorithm_index][nth_run] = indicator_value
        cls.multiple_indicators[indicator_data['name']] = temp_indicator
        return indicator_value

    @staticmethod
    @functools.lru_cache
    def get_cls_by_str(cls_str):
        """
        将data字典中的cls字符串转变成class
        :param:  data.cls形如 algorithms.large.lmoeads.LMOEADS
        :return:
        """
        module_str, obj_name = cls_str.rsplit('.', 1)
        module = importlib.import_module(module_str)
        return getattr(module, obj_name)

    @staticmethod
    def create_3D_list(x, y, z, init_value=None):
        """
        创建指定大小的3D列表
        :param x: X 维度大小
        :param y: Y 维度大小
        :param z: Z 维度大小
        :param init_value: 初始数值，默认为 None
        :return: 3D 列表
        """
        return [[[init_value for _ in range(z)] for _ in range(y)] for _ in range(x)]

    @staticmethod
    def create_2D_list(x, y, init_value=None):
        """
        创建指定大小的3D列表
        :param x: X 维度大小
        :param y: Y 维度大小
        :param z: Z 维度大小
        :param init_value: 初始数值，默认为 None
        :return: 3D 列表
        """
        return [[init_value for _ in range(y)] for _ in range(x)]

    @classmethod
    def return_res(cls, problem_index, algorithm_index, indicator_data):
        temp = cls.multiple_indicators[indicator_data['name']][problem_index][algorithm_index]
        cls.callback({
            'messageType': MessageType.MULTIPLE_TASK_MSG.value,
            'indicator_name': indicator_data['name'],
            'indicator_value': ''.join([f'{temp.mean().item():.3e}', ' (',
                                        f'{temp.std().item():.3e}', ') ']),
            'problem_index': problem_index,
            'algorithm_index': algorithm_index,
        })

    @staticmethod
    def num_of_run_return(algorithm_index, problem_index):
        MsgContext.callback({
            'messageType': MessageType.MULTIPLE_TASK_MSG.value,
            'indicator_name': MsgContext.num_of_run,
            'indicator_value': 1,
            'problem_index': problem_index,
            'algorithm_index': algorithm_index
        })

    @staticmethod
    def get_data(message):
        algorithms_datas = message['uploadAlgorithmsData']
        problem_datas = message['uploadProblemData']
        form_data = message['formData']
        indicator_data = message['indicator_data']
        statistical_type = message['statistical_type']
        return algorithms_datas, problem_datas, form_data, indicator_data, statistical_type

    @staticmethod
    def get_form_data(form_data):
        var_dim = form_data['var_dim']
        obj_dim = form_data['obj_dim']
        pop_size = form_data['pop_size']
        max_fun_eval = form_data['max_fun_eval']
        save_size = form_data['save_size']
        run_size = form_data['run_size']
        return max_fun_eval, obj_dim, pop_size, run_size, save_size, var_dim

    @staticmethod
    def task_complete_return(messageType, tips=None):
        MsgContext.callback({
            'messageType': messageType.value,
            "tips": tips
        })

    @staticmethod
    def multi_statistical_result(statistical_result, best_value_index):
        MsgContext.callback({
            'messageType': MessageType.MULTI_STATISTICAL_RESULT.value,
            'statistical_result': statistical_result,
            "best_value_index": best_value_index.tolist()
        })

    @classmethod
    def statistical_result(cls, algorithms_data_length, problem_datas_length, indicator_data, statistical_type):
        indicator_name = indicator_data['name']
        indicator_class = MsgContext.get_cls_by_str(indicator_data['cls'])
        statistical_result = MsgContext.create_2D_list(problem_datas_length + 1, algorithms_data_length, init_value="")
        best_value_index = torch.zeros(problem_datas_length, algorithms_data_length, dtype=torch.bool)
        for index in range(problem_datas_length):
            temp_indicator = MsgContext.multiple_indicators[indicator_name][index]
            mean_value = temp_indicator.mean(dim=1)
            if indicator_class.isTheSmallerTheBetter():
                best_value_index[index][mean_value.argmin()] = 1
            else:
                best_value_index[index][mean_value.argmax()] = 1
            diff = []
            for temp in temp_indicator[:-1]:
                try:
                    result = cls.statistical_func[statistical_type](temp, temp_indicator[-1])[1] < 0.05
                except ValueError:
                    result = False
                diff.append(result)
            diff = torch.tensor(diff, dtype=torch.bool, device=device)
            for i in range(algorithms_data_length - 1):
                if not diff[i] or mean_value[i] == mean_value[-1]:
                    statistical_result[index][i] = '='
                elif mean_value[i] < mean_value[-1] and indicator_class.isTheSmallerTheBetter() or mean_value[i] > \
                        mean_value[-1] and not indicator_class.isTheSmallerTheBetter():
                    statistical_result[index][i] = '+'
                else:
                    statistical_result[index][i] = '-'
        for j in range(algorithms_data_length - 1):
            num_plus_per = 0
            num_minus_per = 0
            num_equal_per = 0
            for i in range(problem_datas_length):
                if statistical_result[i][j] == '+':
                    num_plus_per += 1
                elif statistical_result[i][j] == '-':
                    num_minus_per += 1
                elif statistical_result[i][j] == '=':
                    num_equal_per += 1
            statistical_result[problem_datas_length][j] = f'{num_plus_per}/{num_minus_per}/{num_equal_per}'

        statistical_result[-1].pop()
        MsgContext.multi_statistical_result(statistical_result, best_value_index)

    @staticmethod
    def save_indicator(algorithm_index, intermediate_pops, nth_run, problem_index):
        indicators_file_name = MsgContext.file_path_list[problem_index][algorithm_index][
            nth_run]
        temp_indicator = {}
        for indicator_name, indicator_value in MsgContext.multiple_indicators.items():
            temp_indicator[indicator_name] = indicator_value[problem_index][algorithm_index][
                nth_run]
        with open(indicators_file_name, 'wb') as my_file:
            pickle.dump([intermediate_pops, temp_indicator], my_file)
